/**********************************************************************
 * file:  sr_arp.c 
 * date:  Thu Oct 18   
 * Contact: yujiali@cs.toronto.edu
 *
 * Description:
 * 
 * This file contains all utilities to handle ARP packets.
 *
 **********************************************************************/

#include <assert.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <pthread.h>
#include <unistd.h>

#include "sr_arp.h"
#include "sr_router.h"
#include "sr_protocol.h"
#include "sr_if.h"
#include "sr_ip.h"
#include "sr_icmp.h"

pthread_mutex_t arp_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t arp_q_mutex;
pthread_mutex_t arp_buf_mutex = PTHREAD_MUTEX_INITIALIZER;

int sr_update_arpent(struct sr_instance*, uint32_t, uint8_t*);
int sr_add_arpent(struct sr_instance*, uint32_t, uint8_t*);

int sr_notify_queue(struct sr_instance* sr, uint32_t ip, uint8_t* ether_addr);
int sr_send_pending_packet(struct sr_instance* sr, struct sr_arp_q* q, 
        uint8_t* ether_addr);

/* only for debugging */
void sr_print_arphdr(struct sr_arphdr* a_hdr);

uint8_t arp_pkt_buf[ARP_PKT_BUF_LEN];

/*--------------------------------------------------------------------- 
 * Method: sr_handle_arp(struct sr_instance* sr, uint8_t* packet, 
 *              unsigned int len)
 * Scope: local
 *
 * This method is called every time an ARP packet is received. The packet
 * buffer is complete with Ethernet headers and ARP packet is inside it.
 *
 * After processing the ARP packet, a 0 is returned if everything goes
 * well, otherwise an error code is returned.
 *
 *---------------------------------------------------------------------*/
int sr_handle_arp(struct sr_instance* sr, 
        uint8_t* packet, 
        unsigned int len,
        char* interface)
{
    /* REQUIRES */
    assert(sr);
    assert(packet);

    struct sr_if*           iface = sr_get_interface(sr, interface);
    struct sr_ethernet_hdr* e_hdr = (struct sr_ethernet_hdr *)packet;
    struct sr_arphdr*       a_hdr = 
        (struct sr_arphdr *)(packet + sizeof(struct sr_ethernet_hdr));
    int merge_flag = 0;

    /* sr_print_arphdr(a_hdr); */
    
    /* check hardware type, should be ethernet */
    if (ntohs(a_hdr->ar_hrd) != ARPHDR_ETHER)
    {
        fprintf(stderr, "Not an Ethernet ARP packet!\n");
        return -1;
    }

    /* check ethernet address length, should be 48 bits */
    if (a_hdr->ar_hln != ETHER_ADDR_LEN) 
    {
        fprintf(stderr, "Ethernet addresses should be %d bytes long, provided" \
                " %d bytes.\n", ETHER_ADDR_LEN, a_hdr->ar_hln);
        return -1;
    }

    /* upper layer protocol should be IP */
    if (ntohs(a_hdr->ar_pro) != ETHERTYPE_IP)
    {
        fprintf(stderr, "ARP not for IP.\n");
        return -1;
    }

    /* check IP address length, should be 32 bits */
    if (a_hdr->ar_pln != IP_ADDR_LEN)
    {
        fprintf(stderr, "IP addresses should be %d bits long, provided" \
                " %d bits.\n", IP_ADDR_LEN, a_hdr->ar_pln);
        return -1;
    }

    merge_flag = sr_update_arpent(sr, a_hdr->ar_sip, a_hdr->ar_sha);

    /* drop the packet if ARP not targeted to us */
    if (a_hdr->ar_tip != iface->ip)
        return 0;
    
    /* create an entry in the ARP table for the sender */
    if (!merge_flag)
        sr_add_arpent(sr, a_hdr->ar_sip, a_hdr->ar_sha);

    switch (ntohs(a_hdr->ar_op))
    {
        case ARP_REPLY:
            /* notify packets in queue to be sent */
            printf("Received an ARP reply message...\n");

            sr_notify_queue(sr, a_hdr->ar_sip, a_hdr->ar_sha);
            break;

        case ARP_REQUEST:

            /* send a reply message, reuse the packet buffer */
            a_hdr->ar_op = htons(ARP_REPLY);
            memcpy(a_hdr->ar_tha, a_hdr->ar_sha, ETHER_ADDR_LEN);
            a_hdr->ar_tip = a_hdr->ar_sip;
            memcpy(a_hdr->ar_sha, iface->addr, ETHER_ADDR_LEN);
            a_hdr->ar_sip = iface->ip;

            /* change the ethernet header as well */
            memcpy(e_hdr->ether_dhost, a_hdr->ar_tha, ETHER_ADDR_LEN);
            memcpy(e_hdr->ether_shost, a_hdr->ar_sha, ETHER_ADDR_LEN);

            /* sr_print_arphdr(a_hdr); */

            sr_send_packet(sr, (uint8_t*)e_hdr, len, interface);
            break;

        default:
            fprintf(stderr, "Unknown ARP operation type: 0x%04x\n", \
                    ntohs(a_hdr->ar_op));
            return -1;
    }

    return 0;
}


/*--------------------------------------------------------------------- 
 * Method: sr_update_arpent(struct sr_instance* sr, uint32_t ip, 
 *              uint8_t* ether_addr)
 * Scope: local
 *
 * This method updates the ARP table entry with the given IP. If there is
 * no entry with the specified IP, return 0. If the entry is successfully 
 * updated, return 1, otherwise return an error code.
 *
 *---------------------------------------------------------------------*/
int sr_update_arpent(struct sr_instance* sr, uint32_t ip, uint8_t* ether_addr)
{
    assert(sr);
    assert(ether_addr);

    int found = 0;
    struct sr_at* at_walker = 0;
    
    /* printf("sr_update_arpent\n"); */

    /* lock the ARP table before reading or changing anything */
    pthread_mutex_lock(&arp_mutex);
    at_walker = sr->arp_table;

    while (at_walker)
    {
        if (at_walker->ip == ip)
        {
            found = 1;
            at_walker->ttl = ARP_ENTRY_INIT_TTL;
            memcpy(at_walker->ether_addr, ether_addr, ETHER_ADDR_LEN);
            break;
        }
        at_walker = at_walker->next;
    }
    pthread_mutex_unlock(&arp_mutex);

    return found;
}

/*--------------------------------------------------------------------- 
 * Method: sr_add_arpent(struct sr_instance* sr, uint32_t ip, 
 *              uint8_t* ether_addr)
 *
 * Scope: local
 *
 * This method add an entry to the ARP table, return 0 if successful, 
 * return an error code otherwise.
 * 
 *---------------------------------------------------------------------*/
int sr_add_arpent(struct sr_instance* sr, uint32_t ip, uint8_t* ether_addr)
{
    assert(sr);
    assert(ether_addr);

    struct sr_at* arp_entry = (struct sr_at *)malloc(sizeof(struct sr_at));
    assert(arp_entry);

    /* printf("sr_add_arpent\n"); */

    /* insert the new entry to the head of the list */
    arp_entry->ip = ip;
    memcpy(arp_entry->ether_addr, ether_addr, ETHER_ADDR_LEN);
    arp_entry->ttl = ARP_ENTRY_INIT_TTL;

    /* lock the ARP table before changing or reading anything */
    pthread_mutex_lock(&arp_mutex);

    arp_entry->next = sr->arp_table;
    sr->arp_table = arp_entry;

    pthread_mutex_unlock(&arp_mutex);

    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_search_arp(struct sr_instance* sr, uint32_t ip, 
 *              uint8_t* ether_addr)
 *
 * Scope: local
 *
 * This method would search the ARP table for a give IP address, if entry
 * found, the ethernet address would be placed in ether_addr array and 
 * return 1, otherwise return 0 and ether_addr is not changed.
 * 
 *---------------------------------------------------------------------*/
int sr_search_arp(struct sr_instance* sr, uint32_t ip, uint8_t* ether_addr)
{
    assert(sr);
    assert(ether_addr);

    /* acquire the lock */
    pthread_mutex_lock(&arp_mutex);

    struct sr_at* at_walker = sr->arp_table;

    while (at_walker)
    {
        if (at_walker->ip == ip)
        {
            memcpy(ether_addr, at_walker->ether_addr, ETHER_ADDR_LEN);

            /* release the lock */
            pthread_mutex_unlock(&arp_mutex);
            return 1;
        }
        at_walker = at_walker->next;
    }

    /* release the lock */
    pthread_mutex_unlock(&arp_mutex);
    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_request_arp(struct sr_instance* sr, uint32_t ip) 
 *
 * Scope: global
 *
 * Send an ARP request for the given IP out from the given interface .
 * Return error code if an error occurs, otherwise return 0.
 * 
 *---------------------------------------------------------------------*/
int sr_request_arp(struct sr_instance* sr, uint32_t ip, struct sr_if* iface)
{
    assert(sr);
    assert(iface);

    struct sr_ethernet_hdr* e_hdr = (struct sr_ethernet_hdr *)arp_pkt_buf;
    struct sr_arphdr* a_hdr = 
        (struct sr_arphdr *)(arp_pkt_buf + sizeof(struct sr_ethernet_hdr));

    /* lock the buffer */
    pthread_mutex_lock(&arp_buf_mutex);

    printf("sr_request_arp\n");

    /* fill in the ethernet header */
    memset(e_hdr->ether_dhost, 0xff, ETHER_ADDR_LEN);
    memcpy(e_hdr->ether_shost, iface->addr, ETHER_ADDR_LEN);
    e_hdr->ether_type = htons(ETHERTYPE_ARP);

    /* fill in the ARP header */
    a_hdr->ar_hrd = htons(ARPHDR_ETHER);
    a_hdr->ar_pro = htons(ETHERTYPE_IP);
    a_hdr->ar_hln = ETHER_ADDR_LEN;
    a_hdr->ar_pln = IP_ADDR_LEN;
    a_hdr->ar_op  = htons(ARP_REQUEST);
    
    memcpy(a_hdr->ar_sha, iface->addr, ETHER_ADDR_LEN);
    a_hdr->ar_sip = iface->ip;
    memset(a_hdr->ar_tha, 0xff, ETHER_ADDR_LEN);
    a_hdr->ar_tip = ip;

    /* sr_print_arphdr(a_hdr); */

    sr_send_packet(sr, (uint8_t*)e_hdr, 
            sizeof(struct sr_ethernet_hdr) + sizeof(struct sr_arphdr),
            iface->name);
    
    pthread_mutex_unlock(&arp_buf_mutex);

    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_add_to_queue(...) 
 *
 * Scope: global
 *
 * Add a IP packet waiting for ARP response to the queue.
 * 
 * ip_hdr is a pointer to the IP packet buffer, complete with IP headers,
 * ip_len is the length of the whole IP packet.
 *---------------------------------------------------------------------*/
int sr_add_to_queue(struct sr_instance* sr, uint32_t ip_gw, 
        struct sr_if* iface, uint8_t* ip_hdr, unsigned int ip_len)
{
    assert(sr);
    assert(iface);
    assert(ip_hdr);

    struct sr_arp_q* aq_walker = 0;
    struct sr_arp_q* new_entry = 0;

    printf("sr_add_to_queue\n");

    new_entry = (struct sr_arp_q *)malloc(sizeof(struct sr_arp_q));
    assert(new_entry);

    new_entry->pkt = (uint8_t *)malloc(sizeof(struct sr_ethernet_hdr) + ip_len);
    assert(new_entry->pkt);

    memcpy(new_entry->pkt + sizeof(struct sr_ethernet_hdr), ip_hdr, ip_len);
    new_entry->ip_gw = ip_gw;
    new_entry->iface = iface;
    new_entry->ttl = ARP_REQUEST_TIMEOUT;
    new_entry->next = 0;

    /* lock the queue */
    pthread_mutex_lock(&arp_q_mutex);

    /* empty queue special case */
    if (!(sr->arp_queue))
    {
        sr->arp_queue = new_entry;
        pthread_mutex_unlock(&arp_q_mutex);
        return 0;
    }

    /* insert the new entry to the end of the queue */
    aq_walker = sr->arp_queue;
    while (aq_walker->next)
        aq_walker = aq_walker->next;

    aq_walker->next = new_entry;

    pthread_mutex_unlock(&arp_q_mutex);
    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_notify_queue(...) 
 *
 * Scope: local
 *
 * Notify the queue when receiving an ARP reply.
 *
 * ip and ether_addr are the IP address and ethernet address requested
 * respectively.
 *---------------------------------------------------------------------*/
int sr_notify_queue(struct sr_instance* sr, uint32_t ip, uint8_t* ether_addr)
{
    assert(sr);
    assert(ether_addr);

    struct sr_arp_q* aq_walker = 0;
    struct sr_arp_q* aq_prev = 0;

    printf("sr_notify_queue\n");

    /* lock the queue */
    pthread_mutex_lock(&arp_q_mutex);

    /* first request special case */
    while (sr->arp_queue)
    {
        aq_walker = sr->arp_queue;
        if (aq_walker->ip_gw == ip)
        {
            /* send the IP packet through ethernet */
            sr_send_pending_packet(sr, aq_walker, ether_addr);

            sr->arp_queue = aq_walker->next;
            free(aq_walker->pkt);
            free(aq_walker);
        }
        else
            break;
    }

    /* empty queue special case */
    if (!(sr->arp_queue))
    {
        pthread_mutex_unlock(&arp_q_mutex);
        return 0;
    }

    aq_prev = aq_walker;
    aq_walker = aq_walker->next;
    while (aq_walker)
    {
        if (aq_walker->ip_gw == ip)
        {
            sr_send_pending_packet(sr, aq_walker, ether_addr);

            aq_walker = aq_walker->next;
            free(aq_prev->next->pkt);
            free(aq_prev->next);
            aq_prev->next = aq_walker;
        }
        else
        {
            aq_prev = aq_walker;
            aq_walker = aq_walker->next;
        }
    }

    pthread_mutex_unlock(&arp_q_mutex);
    return 0;
}


/*--------------------------------------------------------------------- 
 * Method: sr_send_pending_packet(...) 
 *
 * Scope: local
 *
 * Send out a pending IP packet waiting for ARP replies.
 *
 * q is a pointer to the entry in ARP queue, ether_addr is the ethernet
 * address in the ARP reply.
 *---------------------------------------------------------------------*/
int sr_send_pending_packet(struct sr_instance* sr, struct sr_arp_q* q, 
        uint8_t* ether_addr)
{
    assert(sr);
    assert(q);
    assert(ether_addr);

    printf("sr_send_pending_packet\n");


    struct sr_ethernet_hdr* e_hdr = (struct sr_ethernet_hdr *)(q->pkt);
    struct ip* ip_hdr = (struct ip *)(q->pkt + sizeof(struct sr_ethernet_hdr));
    unsigned int ether_len = ntohs(ip_hdr->ip_len) 
        + sizeof(struct sr_ethernet_hdr);

    memcpy(e_hdr->ether_dhost, ether_addr, ETHER_ADDR_LEN);
    memcpy(e_hdr->ether_shost, q->iface->addr, ETHER_ADDR_LEN);
    e_hdr->ether_type = htons(ETHERTYPE_IP);

    /* sr_print_iphdr(ip_hdr); */

    return sr_send_packet(sr, (uint8_t*)e_hdr, ether_len, q->iface->name);
}

/*--------------------------------------------------------------------- 
 * Method: sr_check_arp_queue(void* sr) 
 *
 * Scope: global
 *
 * This function checks the ARP queue and decrement TTL for each entry
 * once per second. Any entry with a zero TTL will be removed from the
 * queue.
 *
 * The input argument should be of type struct sr_instance*.
 *---------------------------------------------------------------------*/
void *sr_check_arp_queue(void* sr_in)
{
    assert(sr_in);

    struct sr_instance* sr = (struct sr_instance *)sr_in;
    struct sr_arp_q* aq_walker = 0;
    struct sr_arp_q* aq_prev = 0;
    struct ip* ip_hdr;

    while (1)
    {
        sleep(1);

        pthread_mutex_lock(&arp_q_mutex);

        /* first request special case */
        while (sr->arp_queue)
        {
            aq_walker = sr->arp_queue;
            aq_walker->ttl -= 1;
            if (aq_walker->ttl == 0)
            {
                /* send ICMP packet */
                ip_hdr = (struct ip *)(aq_walker->pkt 
                        + sizeof(struct sr_ethernet_hdr));

                sr_send_icmp_packet(sr, 
                        ICMP_TYPE_DST_UNREACHABLE, 
                        ICMP_CODE_HOST_UNREACHABLE,
                        ip_addr_to_uint32(ip_hdr->ip_src), 
                        aq_walker->iface->ip, 
                        (uint8_t *)ip_hdr);

                sr->arp_queue = aq_walker->next;
                free(aq_walker->pkt);
                free(aq_walker);
            }
            else
                break;
        }

        /* empty queue */
        if (!(sr->arp_queue))
        {
            pthread_mutex_unlock(&arp_q_mutex);
            continue;
        }

        aq_prev = aq_walker;
        aq_walker = aq_walker->next;

        while (aq_walker)
        {
            aq_walker->ttl -= 1;
            if (aq_walker->ttl == 0)
            {
                /* send ICMP packet */
                ip_hdr = (struct ip *)(aq_walker->pkt 
                        + sizeof(struct sr_ethernet_hdr));

                sr_send_icmp_packet(sr, 
                        ICMP_TYPE_DST_UNREACHABLE, 
                        ICMP_CODE_HOST_UNREACHABLE,
                        ip_addr_to_uint32(ip_hdr->ip_src), 
                        aq_walker->iface->ip, 
                        (uint8_t *)ip_hdr);

                aq_walker = aq_walker->next;
                free(aq_prev->next->pkt);
                free(aq_prev->next);
                aq_prev->next = aq_walker;
            }
            else
            {
                aq_prev = aq_walker;
                aq_walker = aq_walker->next;
            }
        }

        pthread_mutex_unlock(&arp_q_mutex);
    }
    return 0;
}

/*--------------------------------------------------------------------- 
 * Method: sr_check_arp_entries(void* sr) 
 *
 * Scope: global
 *
 * This function checks the ARP table and decrement TTL for each entry
 * once per second. Any entry with a zero TTL will be removed from the
 * table.
 *
 * The input argument should be of type struct sr_instance*.
 *---------------------------------------------------------------------*/
void *sr_check_arp_entries(void* sr_in)
{
    assert(sr_in);

    struct sr_at* at_walker = 0;
    struct sr_at* at_prev = 0;
    struct sr_instance* sr = (struct sr_instance *)sr_in;

    while (1) 
    {
        /* check arp table once per second */
        sleep(1);

        pthread_mutex_lock(&arp_mutex);

        /* empty table special case */
        if (!(sr->arp_table))
        {
            pthread_mutex_unlock(&arp_mutex);
            continue;
        }

        /* first entry special case */
        at_walker = sr->arp_table;
        while (at_walker)
        {
            at_walker->ttl -= 1;
            if (at_walker->ttl == 0)
            {
                at_walker = at_walker->next;
                free(sr->arp_table);
                sr->arp_table = at_walker;
            }
            else
                break;
        }

        /* the table become empty */
        if (!at_walker)
        {
            pthread_mutex_unlock(&arp_mutex);
            continue;
        }

        /* table not empty */
        at_prev = at_walker;
        at_walker = at_walker->next;

        while (at_walker)
        {
            at_walker->ttl -= 1;
            if (at_walker->ttl == 0)
            {
                at_walker = at_walker->next;
                free(at_prev->next);
                at_prev->next = at_walker;
            }
            else
            {
                at_prev = at_walker;
                at_walker = at_walker->next;
            }
        }

        pthread_mutex_unlock(&arp_mutex);
    }

    return 0;
}


/* Print out the ARP header in a readable format */
void sr_print_arphdr(struct sr_arphdr* a_hdr)
{
    assert(a_hdr);

    uint32_t ip;

    printf("ar_hrd: %d\n",      ntohs(a_hdr->ar_hrd));
    printf("ar_pro: 0x%04x\n",  ntohs(a_hdr->ar_pro));
    printf("ar_hln: %d\n",      a_hdr->ar_hln);
    printf("ar_pln: %d\n",      a_hdr->ar_pln);
    printf("ar_op : %d\n",      ntohs(a_hdr->ar_op));
    printf("ar_sha: %02x",      a_hdr->ar_sha[0]);
    for (int i = 1; i < ETHER_ADDR_LEN; i++)
        printf(":%02x", a_hdr->ar_sha[i]);
    printf("\n");

    ip = ntohl(a_hdr->ar_sip);
    printf("ar_sip: %u.%u.%u.%u\n",
            (unsigned short) ((ip & 0xff000000) >> 24),
            (unsigned short) ((ip & 0x00ff0000) >> 16),
            (unsigned short) ((ip & 0x0000ff00) >> 8),
            (unsigned short) ( ip & 0x000000ff));

    printf("ar_tha: %02x",      a_hdr->ar_tha[0]);
    for (int i = 1; i < ETHER_ADDR_LEN; i++)
        printf(":%02x", a_hdr->ar_tha[i]);
    printf("\n");
    
    ip = ntohl(a_hdr->ar_tip);
    printf("ar_tip: %u.%u.%u.%u\n",
            (unsigned short) ((ip & 0xff000000) >> 24),
            (unsigned short) ((ip & 0x00ff0000) >> 16),
            (unsigned short) ((ip & 0x0000ff00) >> 8),
            (unsigned short) ( ip & 0x000000ff));
}

